# -*- coding: utf-8 -*-

"""剑指 Offer II 060. 出现频率最高的 k 个数字
给定一个整数数组 nums 和一个整数 k ，请返回其中出现频率前 k 高的元素。可以按 任意顺序 返回答案。

示例 1:
输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2:
输入: nums = [1], k = 1
输出: [1]

提示：
1 <= nums.length <= 10^5
k 的取值范围是 [1, 数组中不相同的元素的个数]
题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

进阶：所设计算法的时间复杂度 必须 优于 O(n log n) ，其中 n 是数组大小。"""


import heapq

from collections import Counter

class Frequent():
    def __init__(self, val, count) -> None:
        self.val = val
        self.count = count
    
    def __lt__(self, frequent):
        return self.count < frequent.count

    def __eq__(self, frequent):
        return self.count == frequent.count


class Solution:
    """此题需要两步，计数+堆化
    关键在于为 Frequent 提供 __lt__, __eq__ 两个方法，使对象可直接比较"""
    def topKFrequent(self, nums, k: int):
        frequents = [Frequent(val, count) for val, count in Counter(nums).most_common()]
        heap = heapq.nlargest(k, frequents)
        return [frequent.val for frequent in heap]


if __name__ == '__main__':
    so = Solution()
    print(so.topKFrequent(nums = [1,1,1,2,2,3], k = 2))
    print(so.topKFrequent(nums = [1], k = 1))

